import os
import re
from PIL import Image
from tqdm import tqdm

data_root=r'E:\dataset\text_det\result2'
correct_root = os.path.join(data_root, 'correct')
incorrect_root = os.path.join(data_root, 'incorrect')

os.makedirs(correct_root, exist_ok=True)
os.makedirs(incorrect_root, exist_ok=True)

cls_map = {'incorrect':1, 'missing':2}

def load_label(path):
    with open(path, 'r') as f:
        data = f.read().split('\n')
    labels = []
    for line in data:
        if len(line)>0:
            items = re.split(r'\s+', line)
            if len(items)==6:
                cls, x,y,w,h, det_type = items
                x,y,w,h = float(x), float(y), float(w), float(h)
                labels.append([cls_map[det_type.lower()], x-w/2,y-h/2,x+w/2,y+h/2])
            elif len(items)==5:
                cls, x, y, w, h = items
                x, y, w, h = float(x), float(y), float(w), float(h)
                labels.append([int(cls), x - w / 2, y - h / 2, x + w / 2, y + h / 2])
            else:
                print(line)
    return labels

imgs = os.listdir(os.path.join(data_root, 'images'))
for img in tqdm(imgs):
    name = img.split('.')[0]
    txt_path = os.path.join(data_root, 'labels', name+'.txt')
    if not os.path.exists(txt_path):
        continue

    img = Image.open(os.path.join(data_root, 'images', img))
    w, h = img.size

    save_path_incorrect = os.path.join(incorrect_root, name + '_{}.jpg')
    save_path_correct = os.path.join(correct_root, name + '_{}.jpg')

    labels = load_label(txt_path)
    for i, bbox in enumerate(labels):
        cls, x1,y1,x2,y2 = bbox
        if x2>x1 and y2>y1:
            x1,y1,x2,y2 = int(float(x1)*w), int(float(y1)*h), int(float(x2)*w), int(float(y2)*h)
            if cls==0:
                img.crop((x1,y1,x2,y2)).save(save_path_correct.format(i))
            elif cls==1:
                img.crop((x1, y1, x2, y2)).save(save_path_incorrect.format(i))
